import numpy as np
import random
from copy import deepcopy
from vrp_model import VRPTWInstance, Customer, Vehicle


class Chromosome:
    """染色体结构体"""

    def __init__(self, genes, cost, routes):
        self.genes = genes  # 原始基因序列
        self.cost = cost  # 总成本（用于比较）
        self.routes = routes  # 解码路径

    def __lt__(self, other):
        """定义比较运算符用于排序"""
        return self.cost < other.cost


class EnhancedGeneticAlgorithm:
    def __init__(self, instance, pop_size=50, min_pop_size=30,
                 max_generations=100, crossover_rate=0.8,
                 mutation_rate=0.1, T=10):
        self.instance = instance          # 问题实例
        self.pop_size = pop_size          # 初始种群大小
        self.min_pop_size = min_pop_size  # 最小种群大小
        self.max_generations = max_generations  # 最大迭代次数
        self.crossover_rate = crossover_rate    # 交叉概率
        self.mutation_rate = mutation_rate      # 变异概率
        self.T = T                        # 种群调整周期
        self.current_pop_size = pop_size  # 当前种群大小
        self.generation = 0               # 当前代数

    def initialize_population(self):
        customers = list(range(1, len(self.instance.customers)))
        return [random.sample(customers, len(customers)) for _ in range(self.pop_size)]

    def calculate_fitness(self, chromo):
        key = tuple(chromo)
        if key in self.instance._fitness_cache:
            return self.instance._fitness_cache[key]

        routes = []
        current_route = [0]
        total_cost = 0
        current_load = 0
        current_volume = 0
        current_time = 0
        vehicle_used = 0

        for customer_id in chromo:
            customer = self.instance.customers[customer_id]

            # 检查换车条件
            if (current_load + customer.demand > self.instance.vehicle.capacity or
                    current_volume + customer.volume > self.instance.vehicle.volume):
                # 计算返回成本
                return_distance = self.instance.distance_matrix[current_route[-1]][0]
                return_time = self.instance.time_matrix[current_route[-1]][0]
                total_cost += return_distance * self.instance.vehicle.distance_cost
                total_cost += return_time * self.instance.vehicle.time_cost
                total_cost += self.instance.vehicle.fixed_cost
                vehicle_used += 1

                # 记录路径
                current_route.append(0)
                routes.append(current_route)

                # 新车初始化
                current_route = [0]
                current_load = 0
                current_volume = 0
                current_time = 0

            # 添加客户到当前路径
            from_node = current_route[-1]
            distance = self.instance.distance_matrix[from_node][customer_id]
            travel_time = self.instance.time_matrix[from_node][customer_id]

            # 时间处理
            arrival_time = current_time + travel_time
            wait_time = max(customer.early - arrival_time, 0)
            time_violation = max(arrival_time - customer.late, 0)

            # 成本累计
            total_cost += distance * self.instance.vehicle.distance_cost
            total_cost += (travel_time + wait_time) * self.instance.vehicle.time_cost
            total_cost += time_violation * self.instance.vehicle.time_penalty

            # 更新状态
            current_route.append(customer_id)
            current_load += customer.demand
            current_volume += customer.volume
            current_time = arrival_time + wait_time + 60

        # 处理最后路径
        if len(current_route) > 1:
            return_distance = self.instance.distance_matrix[current_route[-1]][0]
            return_time = self.instance.time_matrix[current_route[-1]][0]
            total_cost += return_distance * self.instance.vehicle.distance_cost
            total_cost += return_time * self.instance.vehicle.time_cost
            total_cost += self.instance.vehicle.fixed_cost
            vehicle_used += 1
            current_route.append(0)
            routes.append(current_route)

        # 超载惩罚
        overload_penalty = 0
        for route in routes:
            total_demand = sum(self.instance.customers[n].demand for n in route[1:-1])
            if total_demand > self.instance.vehicle.capacity:
                overload_penalty += self.instance.vehicle.overload_penalty
        total_cost += overload_penalty

        # 缓存结果
        chromo_obj = Chromosome(chromo, total_cost, routes)
        self.instance._fitness_cache[key] = chromo_obj
        return chromo_obj

    def cyclic_population_adjustment(self, gen):
        cycle = self.max_generations // 5
        if gen % cycle == 0 and gen != 0:
            self.current_pop_size = max(self.min_pop_size, int(self.current_pop_size * 0.8))

    def pmx_crossover(self, parent1, parent2):
        size = len(parent1)
        p1, p2 = random.sample(range(size), 2)
        start, end = min(p1, p2), max(p1, p2)

        child1 = parent1[:start] + parent2[start:end] + parent1[end:]
        child2 = parent2[:start] + parent1[start:end] + parent2[end:]

        # 修复重复基因
        for child in [child1, child2]:
            seen = set()
            duplicates = []
            for i, gene in enumerate(child):
                if gene in seen:
                    duplicates.append(i)
                else:
                    seen.add(gene)
            missing = set(parent1) - seen
            for idx in duplicates:
                child[idx] = missing.pop()
        return child1, child2

    '''两点交换变异'''
    def adaptive_mutation(self, chromo):
        if random.random() < self.mutation_rate:
            idx1, idx2 = random.sample(range(len(chromo)), 2)
            chromo[idx1], chromo[idx2] = chromo[idx2], chromo[idx1]
        return chromo

    def defense_strategy_enhancement(self, parent, population):
        fitness_values = [self.calculate_fitness(c).cost for c in population]
        median_fitness = np.median(fitness_values)
        std_fitness = np.std(fitness_values)
        current_fitness = self.calculate_fitness(parent).cost
        return current_fitness > (median_fitness + 0.5 * std_fitness)

    def cpo_selection(self, population, fitness_values):
        sorted_indices = np.argsort(fitness_values)
        population_size = len(population)

        selection_pressure = 0.8 - 0.6 * (self.generation / self.max_generations)
        ranks = np.arange(1, population_size + 1)
        probabilities = 1 / (
                    population_size * (selection_pressure + (1 - selection_pressure) * ranks / population_size))
        probabilities /= probabilities.sum()

        selected_indices = np.random.choice(
            population_size,
            size=population_size,
            replace=True,
            p=probabilities
        )
        return [population[i] for i in selected_indices]

    def solve(self):
        population = self.initialize_population()
        best_solution = None
        best_fitness = float('inf')
        history = []

        for gen in range(self.max_generations):
            self.cyclic_population_adjustment(gen)
            population = population[:self.current_pop_size]

            # 适应度计算
            fitness_objects = [self.calculate_fitness(c) for c in population]
            fitness_values = [f.cost for f in fitness_objects]

            # 选择
            selected = self.cpo_selection(population, fitness_values)

            # 交叉与变异
            new_population = []
            print("种群大小：{}".format(self.current_pop_size))
            while len(new_population) < self.current_pop_size:
                parents = random.sample(selected, 2)
                if random.random() < self.crossover_rate:
                    child1, child2 = self.pmx_crossover(parents[0], parents[1])
                else:
                    child1, child2 = deepcopy(parents[0]), deepcopy(parents[1])

                new_population.extend([
                    self.adaptive_mutation(child1),
                    self.adaptive_mutation(child2)
                ])

            population = new_population[:self.current_pop_size]
            current_best = min(fitness_objects, key=lambda x: x.cost)

            if current_best.cost < best_fitness:
                best_solution = current_best
                best_fitness = current_best.cost

            history.append(best_fitness)
            print(f"Generation {gen + 1}: Best Cost={best_fitness:.2f} Routes={len(best_solution.routes)}")
            self.instance.clear_cache()
            self.generation += 1

        return best_solution, history